from diffusers import AutoPipelineForText2Image
import torch
import json
from http.server import BaseHTTPRequestHandler, HTTPServer

pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
pipe.to("cuda")

class MyHandler(BaseHTTPRequestHandler):

    def _handle_request(self, client_data: dict) -> str:
        torch.manual_seed(client_data['seed'])

        image = pipe(
            prompt=client_data['prompt'],
            num_inference_steps=client_data['steps'],
            guidance_scale=0.0,
            height=client_data['height'],
            width=client_data['width'],
        ).images[0]
        image.save(client_data['output'], 'PNG', optimize=True)

    def do_POST(self):
        # 获取POST推送过来的内容
        content_length = int(self.headers['Content-Length'])
        if content_length > 1024 * 1024:
            return

        post_data = self.rfile.read(content_length)
        try:
            client_data = json.loads(post_data.decode(encoding='utf8', errors='replace'))
            print(json.dumps(client_data, ensure_ascii=False))
        except json.decoder.JSONDecodeError as e:
            print(e)

        self._handle_request(client_data)

        # 响应文本内容：ok
        self.send_response(200)
        self.send_header('Content-type', 'text/plain')
        self.end_headers()
        self.wfile.write(b'ok')

ser = HTTPServer(('127.0.0.1', 8977), MyHandler) 
ser.serve_forever()
